
"""
This is part of Human Robot Dynamics-Control, By Tourists.

Based on "OpenLoong Dynamics Control"
Transfer from mujoco to ISAAC SIM by IsaacLab

"""
import numpy as np
import glfw
from OpenGL.GL import *
import mujoco
#pragma once
#include <mujoco/mujoco.h>
#include <GLFW/glfw3.h>
#include <string>
#include <memory>

class ButtonState:
    def __init__(self):
        self.key_w = False
        self.key_s = False
        self.key_a = False
        self.key_d = False
        self.key_h = False
        self.key_j = False
        self.key_space = False




class UIctr:
    def __init__(self, modelIn, dataIn):
        # self.window = glfw.create_window(1280, 720, "MuJoCo Viewer", None, None)
        self.window = None
        ##
        self.buttonRead = ButtonState()

        ## mouse interaction
        self.button_left = False
        self.button_middle = False
        self.button_right = False

        self.runSim = True
        self.isContinuous = True
        self.lastx = 0.0
        self.lasty = 0.0
        
        self.mj_model = modelIn
        self.mj_data = dataIn
        self.cam = mujoco.MjvCamera()
        self.opt = mujoco.MjvOption()
        self.scn = mujoco.MjvScene()
        self.con = mujoco.MjrContext()
        
        self.width = 1200
        self.height = 800
        
        self.save_video = False
        self.isTrack = False
        # Equivalent to unsigned char* image_rgb_;
        self.image_rgb = np.zeros((0, 0, 3), dtype=np.uint8)
        # Equivalent to float* image_depth_;
        self.image_depth = np.zeros((0, 0), dtype=np.float32)
        
        # self.file = open('filename.txt', 'r')
        self.file = None
        
    def iniGLFW(self):
        if not glfw.init():
            raise Exception("Could not initialize GLFW")
        # self.window = glfw.create_window(1280, 720, "MuJoCo Viewer", None, None)

    def createWindow(self, windowTitle, saveVideo):
        self.window = glfw.create_window(self.width, self.height, windowTitle, None, None)
        if not self.window:
            glfw.terminate()
            raise Exception("Could not create window")

        glfw.make_context_current(self.window)
        glfw.swap_interval(1)

        # Set up camera
        # self.cam = mujoco.MjvCamera()
        mujoco.mjv_defaultFreeCamera(self.mj_model, self.cam)
        arr_view = [150, -16, 3, 0, 0.000000, 1.00000]
        self.cam.azimuth = arr_view[0]
        self.cam.elevation = arr_view[1]
        self.cam.distance = arr_view[2]
        self.cam.lookat[:] = arr_view[3:6]

        if self.isTrack:
            self.cam.lookat[2] += 0.8
            self.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
            self.cam.trackbodyid = 1
        else:
            self.cam.type = mujoco.mjtCamera.mjCAMERA_FREE

        # mujoco.mjv_defaultOption(self.opt)
        # mujoco.mjv_defaultScene(self.scn)
        # mujoco.mjr_defaultContext(self.con)
        self.scn = mujoco.MjvScene(self.mj_model, 2000);                ## space for 2000 objects
        # self.con = mujoco.MjrContext()(self.mj_model, self.con, mujoco.mjtFontScale.mjFONTSCALE_150);   
        self.con = mujoco.MjrContext(self.mj_model, mujoco.mjtFontScale.mjFONTSCALE_150.value)   ## model-specific context
        # mujoco.mjv_moveCamera(self.mj_model, mujoco.mjtMouse.mjMOUSE_ROTATE_H, 0.0, 0.0, self.scn, self.cam)
        
        # Set callbacks
        glfw.set_window_user_pointer(self.window, self)
        glfw.set_window_close_callback(self.window, self.window_close_callback)
        glfw.set_key_callback(self.window, self.keyboard)
        glfw.set_cursor_pos_callback(self.window, self.mouse_move)
        glfw.set_mouse_button_callback(self.window, self.mouse_button)
        glfw.set_scroll_callback(self.window, self.scroll)

        self.save_video = saveVideo
        if self.save_video:
            self.image_rgb_ = np.zeros((self.height, self.width, 3), dtype=np.uint8)
            self.image_depth_ = np.zeros((self.height, self.width), dtype=np.float32)
            self.file = open("../record/rgbRec.out", "wb")

    def updateScene(self):
        if not self.isContinuous:
            self.runSim = False

        self.buttonRead = ButtonState()
        
        glfw.make_context_current(self.window)
        width, height = glfw.get_framebuffer_size(self.window)
        viewport = mujoco.MjrRect(0, 0, width, height)

        # self.scn = mujoco.MjvScene(self.mj_model, maxgeom=100)
        
        mujoco.mjv_updateScene(self.mj_model, self.mj_data, self.opt, None, self.cam,
                               mujoco.mjtCatBit.mjCAT_ALL, self.scn)
        mujoco.mjr_render(viewport, self.scn, self.con)

        timeStr = f"Time: {self.mj_data.time:.3f}"
        mujoco.mjr_overlay(mujoco.mjtFont.mjFONT_NORMAL, mujoco.mjtGridPos.mjGRID_TOPRIGHT,
                           viewport, timeStr.encode(), None, self.con)

        glfw.swap_buffers(self.window)
        glfw.poll_events()

        if self.save_video:
            mujoco.mjr_readPixels(self.image_rgb_, self.image_depth_, viewport, self.con)
            self.file.write(self.image_rgb_.tobytes())

    def keyboard(self, window, key, scancode, act, mods):
        if act == glfw.PRESS and key == glfw.KEY_BACKSPACE:
            mujoco.mj_resetData(self.mj_model, self.mj_data)
            mujoco.mj_forward(self.mj_model, self.mj_data)

        if act == glfw.RELEASE:
            if key == glfw.KEY_1:
                self.runSim = not self.runSim
                self.isContinuous = True
            elif key == glfw.KEY_2:
                self.runSim = True
                self.isContinuous = False
            elif key == glfw.KEY_W:
                self.buttonRead.key_w = True
            elif key == glfw.KEY_A:
                self.buttonRead.key_a = True
            elif key == glfw.KEY_S:
                self.buttonRead.key_s = True
            elif key == glfw.KEY_D:
                self.buttonRead.key_d = True
            elif key == glfw.KEY_H:
                self.buttonRead.key_h = True
            elif key == glfw.KEY_J:
                self.buttonRead.key_j = True
            elif key == glfw.KEY_SPACE:
                self.buttonRead.key_space = True

    def mouse_button(self, window, button, act, mods):
        self.button_left = glfw.get_mouse_button(window, glfw.MOUSE_BUTTON_LEFT) == glfw.PRESS
        self.button_middle = glfw.get_mouse_button(window, glfw.MOUSE_BUTTON_MIDDLE) == glfw.PRESS
        self.button_right = glfw.get_mouse_button(window, glfw.MOUSE_BUTTON_RIGHT) == glfw.PRESS
        self.lastx, self.lasty = glfw.get_cursor_pos(window)

    def mouse_move(self, window, xpos, ypos):
        if not (self.button_left or self.button_middle or self.button_right):
            return

        dx = xpos - self.lastx
        dy = ypos - self.lasty
        self.lastx = xpos
        self.lasty = ypos

        width, height = glfw.get_window_size(window)

        mod_shift = glfw.get_key(window, glfw.KEY_LEFT_SHIFT) == glfw.PRESS or \
                    glfw.get_key(window, glfw.KEY_RIGHT_SHIFT) == glfw.PRESS

        if self.button_right:
            action = mujoco.mjtMouse.mjMOUSE_MOVE_H if mod_shift else mujoco.mjtMouse.mjMOUSE_MOVE_V
        elif self.button_left:
            action = mujoco.mjtMouse.mjMOUSE_ROTATE_H if mod_shift else mujoco.mjtMouse.mjMOUSE_ROTATE_V
        else:
            action = mujoco.mjtMouse.mjMOUSE_ZOOM

        mujoco.mjv_moveCamera(self.mj_model, action, dx/height, dy/height, self.scn, self.cam)

    def scroll(self, window, xoffset, yoffset):
        mujoco.mjv_moveCamera(self.mj_model, mujoco.mjtMouse.mjMOUSE_ZOOM, 0, 0.05 * yoffset, self.scn, self.cam)

    def window_close_callback(self, window):
        self.close()

    def close(self):
        if self.save_video:
            self.file.close()
        glfw.terminate()

    def enableTracking(self):
        self.isTrack = True

    def getButtonState(self):
        tmp = self.buttonRead
        self.buttonRead = ButtonState()
        return tmp